﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Xml;
using System.IO;
using System.Diagnostics;

namespace XmlDSigEx
{
    public class XmlSigner : ICreateReference
    {
        public const string NS_XMLDSIG = "http://www.w3.org/2000/09/xmldsig#";
        private List<Reference> _references = new List<Reference>();
        private List<DataObject> _dataObjects = new List<DataObject>();
        private XmlDocument _context;
        private XmlDocument _output = new XmlDocument() { PreserveWhitespace = true };
        private ISignatureAlgorithm _signatureAlgorithm;
       

        public string CanonicalisationMethod { get; set; }
        public string SignatureMethod { get; set; }
        public string ID { get; set; }

        public XmlSigner(XmlDocument context)
        {
            _context = context;
            InitialiseAlgorithms();
        }

        public XmlSigner()
        {
            _context = _output;
            InitialiseAlgorithms();
        }

        private void InitialiseAlgorithms()
        {
            CanonicalisationMethod = ExcXmlC14NCanonicalisation.URI_EXC_C14N;
            SignatureMethod = RsaAlgorithm.URI_RSA_SHA256;
        }

        public Reference CreateReference()
        {
            Reference reference = new Reference();
            _references.Add(reference);
            return reference;
        }

        public XmlElement GetXml(ISignatureAlgorithm signatureAlgorithm)
        {
            using (Stream buffer = new MemoryStream())
            using (XmlWriter writer = XmlWriter.Create(buffer))
            {
                writer.WriteStartElement("Signature", NS_XMLDSIG);                
                if (!string.IsNullOrEmpty(ID))
                    writer.WriteAttributeString("Id", ID);
                WriteSignedInfo(writer);
                writer.WriteElementString("SignatureValue", NS_XMLDSIG, null);
                if (signatureAlgorithm.WillWriteKeyInfo)
                {
                    writer.WriteStartElement("KeyInfo", XmlSigner.NS_XMLDSIG);
                    signatureAlgorithm.WriteKeyInfo(writer);
                    writer.WriteEndElement();
                }
                
                ProcessDataObjects(writer);

                writer.WriteEndElement();
                writer.Flush();
                buffer.Position = 0;                
                _output.Load(XmlReader.Create(buffer));
            }
            
            ProcessReferences(_output);
            ProcessSignature(_output, signatureAlgorithm);
            return _output.DocumentElement;
        }

        private void ProcessDataObjects(XmlWriter writer)
        {
            foreach (DataObject dataObject in _dataObjects)
            {
                writer.WriteStartElement("Object", NS_XMLDSIG);
                if (dataObject.ID != null)
                    writer.WriteAttributeString("Id", dataObject.ID);
                if(dataObject.Creating != null)
                    dataObject.Creating(writer);
                writer.WriteEndElement();
            }
        }

        private void ProcessSignature(XmlDocument output, ISignatureAlgorithm signatureAlgorithm)
        {

            byte[] signature = CalculateSignatureValue(output.DocumentElement.FirstChild as XmlElement, signatureAlgorithm);
            output.DocumentElement.SelectSingleNode("//*[local-name() = 'SignatureValue']").InnerText = Convert.ToBase64String(signature);
        }

        private void ProcessReferences(XmlDocument output)
        {
            XmlNamespaceManager nsMgr = new XmlNamespaceManager(output.NameTable);
            nsMgr.AddNamespace("ds", XmlSigner.NS_XMLDSIG);
            XmlNodeList referenceElements = output.SelectNodes("/ds:Signature/ds:SignedInfo/ds:Reference", nsMgr);
            for (int i = 0; i < referenceElements.Count; i++)
            {
                if (!_references[i].Precomputed)
                    _references[i].ProcessDigest(_context, referenceElements[i] as XmlElement);
            }
        }         
    
        private  byte[] CalculateSignatureValue(XmlElement signedInfo, ISignatureAlgorithm signatureAlgorithm)
        {
            Stream signatureStream = GetSignedInfoStream(signedInfo);
            return signatureAlgorithm.SignData(signatureStream, SignatureMethod);
        }

        private Stream GetSignedInfoStream(XmlElement signedInfo)
        {
            XmlNodeList nodesToSign = signedInfo.SelectNodes("(.//. | .//@* | .//namespace::*)");
            ICanonicalisationAlgorithm canonicalisation = CanonicalisationFactory.GetCanonicalisation(CanonicalisationMethod);
            Stream s = canonicalisation.ProcessNodeSet(nodesToSign, false);
          
            return s;
        }

        private void WriteSignedInfo(XmlWriter writer)
        {
            writer.WriteStartElement("SignedInfo", NS_XMLDSIG);

            writer.WriteStartElement("CanonicalizationMethod", NS_XMLDSIG);
            writer.WriteAttributeString("Algorithm", CanonicalisationMethod);
            writer.WriteEndElement();

            writer.WriteStartElement("SignatureMethod", NS_XMLDSIG);
            writer.WriteAttributeString("Algorithm", SignatureMethod);
            writer.WriteEndElement();

            foreach (Reference reference in _references)
            {
                reference.WriteXml(writer);
            }
            writer.WriteEndElement();

        }


        public bool ValidateSignature(XmlElement signature)
        {
            ISignatureAlgorithm algorithm = CreateAlgorithmFromSignature(signature);
            return ValidateSignature(signature, algorithm);
        }

        private ISignatureAlgorithm CreateAlgorithmFromSignature(XmlElement keyInfoElement)
        {
            return SignatureAlgorithmFactory.GetSignatureAlgorithm(keyInfoElement);
        }

        public bool ValidateSignature(XmlElement signature, ISignatureAlgorithm signatureAlgorithm)
        {
            _signatureAlgorithm = signatureAlgorithm;
            if (signature.LocalName != "Signature" || signature.NamespaceURI != NS_XMLDSIG)
                throw new ArgumentException("Element was not a signature.");

            XmlNamespaceManager manager = new XmlNamespaceManager(signature.OwnerDocument.NameTable);
            manager.AddNamespace("ds", NS_XMLDSIG);

            XmlNodeList nodes = signature.SelectNodes("./ds:SignedInfo/ds:Reference",manager);
            foreach (XmlElement referenceElement in nodes)
            {
                Reference reference = CreateReference();
                if (referenceElement.HasAttribute("URI"))
                    reference.Uri = referenceElement.Attributes["URI"].Value;

                XmlElement digestMethod = referenceElement.SelectSingleNode("./ds:DigestMethod", manager) as XmlElement;
                if(digestMethod != null)
                    reference.DigestMethod = digestMethod.Attributes["Algorithm"].Value;
                XmlElement digestValue = referenceElement.SelectSingleNode("./ds:DigestValue", manager) as XmlElement;
                if (digestMethod != null)
                    reference.DigestValue = Convert.FromBase64String(digestValue.InnerText);

                foreach(XmlElement transform in referenceElement.SelectNodes("./ds:Transforms/ds:Transform", manager))
                {
                    reference.AddTransform(transform.Attributes["Algorithm"].Value, transform);
                }

                bool result = reference.Check(_context, referenceElement);
                if (!result)
                    return false;
            }

            XmlElement signatureValueElement = signature.SelectSingleNode("./ds:SignatureValue", manager) as XmlElement;
            XmlElement signatureMethodElement = signature.SelectSingleNode("./ds:SignedInfo/ds:SignatureMethod", manager) as XmlElement;
            CanonicalisationMethod = (signature.SelectSingleNode("./ds:SignedInfo/ds:CanonicalizationMethod", manager) as XmlElement).Attributes["Algorithm"].Value;

            Stream signedInfoStream = GetSignedInfoStream(signature.FirstChild as XmlElement);

            SignatureValue = Convert.FromBase64String(signatureValueElement.InnerText);

            return signatureAlgorithm.VerifySignature(signedInfoStream, 
                SignatureValue,
                signatureMethodElement.Attributes["Algorithm"].Value); 
        }

        public IDataObject CreateDataObject()
        {
            return CreateDataObject(null);
        }

        public IDataObject CreateDataObject(string id)
        {
            DataObject dataObject = new DataObject() { ID = id };            
            _dataObjects.Add(dataObject);
            return dataObject;
        }

        public IList<Reference> References
        {
            get
            {
                return _references.AsReadOnly();
            }
        }

        public ISignatureAlgorithm GetSignatureAlgorithm()
        {
            return _signatureAlgorithm;
        }

        public byte[] SignatureValue
        {
            get;
            private set;
        }
    }
}
